package com.dahuyou.javaagent.and.javassist;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

/**
 * 自定义字节码文件转换器
 */
public class MyWasteTimeTransformer implements ClassFileTransformer {
    /*
    要增强的类的集合
     */
    private static final Set<String> classNameSet = new HashSet<>();

    static {
        classNameSet.add("com.dahuyou.javaagent.and.javassist.MyApiTest");
    }

    /**
     * 转换方法
     *
     * @param loader
     * @param className           类名称 注意是/格式的
     * @param classBeingRedefined
     * @param protectionDomain
     * @param classfileBuffer     原始的解码文件，即需要插桩的程序对应的字节码文件
     * @return
     * @throws IllegalClassFormatException
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            String fullQualifiedClazzName = className.replace("/", ".");
            if (!classNameSet.contains(fullQualifiedClazzName)) {
                return null;
            }
            System.out.println("start transfer class: " + fullQualifiedClazzName);

            // javassist 增强开始啦
            ClassPool pool = ClassPool.getDefault();
            CtClass ctClass = pool.get(fullQualifiedClazzName);
            // 获取所有的方法，并增加方法耗时代码
            CtBehavior[] declaredBehaviors = ctClass.getDeclaredBehaviors();
            for (CtBehavior method : declaredBehaviors) {
                // 增强
                enhanceMethod(method);
            }
            // 返回插桩增强后的字节码
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }


//        return new byte[0];
        return null;
    }

    /*
    插桩方法执行耗时
     */
    private void enhanceMethod(CtBehavior method) throws Exception {
        String methodName = "";
        if (method.isEmpty() || "main".equalsIgnoreCase(methodName = method.getName())) {
            return;
        }
        final StringBuilder source = new StringBuilder();
        // 前置增强: 打入时间戳
        // 保留原有的代码处理逻辑
        source.append("{")
                .append("long start = System.nanoTime();\n") //前置增强: 打入时间戳
                .append("$_ = $proceed($$);\n") //调用原有代码， 类似于method();($$)表示所有的参数（这个就高级了啊！）
                .append("System.out.print(\"method:[")
                .append(methodName).append("]\");").append("\n")
                .append("System.out.println(\" cost:[\" +(System.nanoTime() - start)+ \"ns]\");") // 后置增强，计算输出方法执行耗时
                .append("}");
        // 方法编辑
        ExprEditor editor = new ExprEditor() {
            @Override
            public void edit(MethodCall methodCall) throws
                    CannotCompileException {
                methodCall.replace(source.toString());
            }
        };
        method.instrument(editor);
    }
}
